from cProfile import label
from turtle import forward
import torch
import torch.nn as nn 
import torch.nn.functional as  F

class SCEloss(nn.Module):
    def __init__(self, num_classes=5, alpha=1, beta=1):
        super().__init__()
        self.num_classes = num_classes
        self.alpha = alpha
        self.beta = beta
        self.cross_entropy = nn.CrossEntropyLoss()

    def forward(self, pred, labels):
        ce = self.cross_entropy(pred, labels)

        pred = F.softmax(pred, dim=1)
        pred = torch.clamp(pred, min=1e-7, max=1.0)
        label_one_hot = F.one_hot(labels, self.num_classes).float().to(pred.device)
        label_one_hot = torch.clamp(label_one_hot, min=1e-4, max=1.0) # A=-4
        rce = (-1 * torch.sum(pred * torch.log(label_one_hot), dim=1))

        loss = self.alpha * ce + self.beta * rce.mean()
        return loss